#!/usr/bin/env python3
# config.py
import argparse
import os
import torch

def parse_args():
    """解析命令行参数，生成配置对象args"""
    parser = argparse.ArgumentParser(description='ESRGAN超分辨率训练配置')
    
    # 基本训练参数
    parser.add_argument('--epochs', type=int, default=100, help='训练总轮数')
    parser.add_argument('--batch_size', type=int, default=16, help='批次大小')
    #parser.add_argument('--lr', type=float, default=0.0002, help='初始学习率')
    parser.add_argument('--g_lr', type=float, default=1e-6, help='生成器学习率')
    parser.add_argument('--d_lr', type=float, default=1e-7, help='判别器学习率')
    parser.add_argument('--lr_decay', type=float, default=0.95, help='学习率衰减系数')
    parser.add_argument('--patch_size', type=int, default=128, 
                       help='Training patch size')
    
    # 模型参数
    parser.add_argument('--scale_factor', type=int, default=4, help='超分辨率缩放倍数')
    parser.add_argument('--num_block', type=int, default=23, help='RRDB块数量')
    parser.add_argument('--num_feat', type=int, default=64, help='特征图数量')
    parser.add_argument('--num_grow_ch', type=int, default=32, 
                       help='Number of growth channels in RDB')
    
    # 数据参数
    parser.add_argument('--dataset_path', type=str, default='./data', help='数据集根目录')
    parser.add_argument('--train_ratio', type=float, default=0.8, help='训练集比例')
    parser.add_argument('--img_size', type=int, default=128, help='输入图像尺寸')
    parser.add_argument('--result_dir', type=str, default='./results', 
                       help='Directory to save results')
    
    # 训练设置
    parser.add_argument('--resume', type=str, default=None, 
                       help='Checkpoint path to resume from')
    parser.add_argument('--checkpoint_interval', type=int, default=100, 
                       help='Save checkpoint every N epochs')
    parser.add_argument('--log_interval', type=int, default=10, 
                       help='Log training status every N batches')
    parser.add_argument('--val_interval', type=int, default=50, 
                       help='Run validation every N epochs')
    
    # 设备与保存参数
    parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu', 
                       help='训练设备（cuda或cpu）')
    parser.add_argument('--checkpoint_dir', type=str, default='./checkpoints', help='模型保存目录')
    parser.add_argument('--log_dir', type=str, default='./logs', help='日志保存目录')
    parser.add_argument('--save_freq', type=int, default=10, help='模型保存频率（轮数）')
    parser.add_argument('--num_workers', type=int, default=4, 
                       help='Number of workers for data loading')
    parser.add_argument('--pin_memory', action='store_true', default=True,
                       help='Use pinned memory for data loading')
    
    # 实验设置
    parser.add_argument('--seed', type=int, default=42, 
                       help='Random seed for reproducibility')
    parser.add_argument('--use_amp', action='store_true', default=False,
                       help='Use automatic mixed precision')
    
    # 添加数据集自动下载参数
    parser.add_argument('--download_datasets', action='store_true', default=False,
                       help='是否自动下载并准备数据集（DIV2K和Flickr2K）')
    
    args = parser.parse_args()
    
    # 创建必要的目录
    os.makedirs(args.checkpoint_dir, exist_ok=True)
    os.makedirs(args.log_dir, exist_ok=True)
    os.makedirs(args.result_dir, exist_ok=True)
    os.makedirs(args.dataset_path, exist_ok=True)
    
    return args

# 生成全局可导入的args对象
args = parse_args()

if __name__ == '__main__':
    test_args = parse_args()
    print("Configuration:")
    for key, value in vars(test_args).items():
        print(f"  {key}: {value}")